# encoding=utf-8
"""BatchPredictService:
    - 将云端上的代码进行批量处理
    - 处理方案参照别人留下来的

    - 旧代码是为封装的代码
    - 新代码将会放弃model_path这个选项(会在app2.py中进行初始化，单例设计模式, 后续将保持于内存中)，
      除非内存溢出，否则在app2.py中初始化完毕后，后续用到可以直接传入 ""
      然后此主流程主要进行逻辑判断, 然后创建相应的对象并送入一个抽象的PredictHandler就行
"""
import json
import os

import flask
from PIL import Image
from StillGAN_API import *
from stage1 import *
from stage2 import predict as grading
from multiprocessing.dummy import Pool as ThreadPool

from seg_system.multibatch.annotation.MultiBatchAnnotation import *
from seg_system import multibatch
from seg_common.logging import ConsoleService
from seg_system.enhance import StillganBatchFactory
from seg_system.enhance import StillganService
from seg_system import ApplicationConfig
from seg_system.segmentation import NerveFormerBatchFactory
from seg_system.segmentation import NerveFormerService
from seg_system.segmentation import CSNetService
from seg_system.segmentation import CSNetBatchFactory
from seg_system.segmentation import NerveCeilBatchFactory, CSNetCeilBatchFactory
from seg_system.segmentation import CSNetCeilService, NerveFormerCeilService
from seg_system.multibatch.service.IPredictHandler import IPredictHandler
from seg_system.LongTaskStatusUpload import MultiStatusUploadImpl
from seg_system.LongTaskStatusUpload import IStatusUpload


""""-------------我也不知道谁写的，我只是一个搬运工----------"""


def get_roi(img, seg):
    """
    First step: obtain the roi
    :param img: numpy array
    :param seg: numpy array
    :return:
    """
    heatmap_on_image, roi, only_heatmap_on_seg, nerve_gradient_on_img = generate_heatmap(
        img, seg)
    # roi will be used in stage2,
    # the only_heatmap_on_seg, nerve_gradient_on_img are used to visualize
    return heatmap_on_image, roi, only_heatmap_on_seg, nerve_gradient_on_img


def tortuosity_grading(img, seg):
    """
    Choose one to show/visualize as you desired
    Second step: obtain the grading results
    :param img: numpy array
    :param seg: numpy array
    :return: 4 items
        "only_heatmap_on_seg": the heatmap on segmentation
        "nerve_gradient_on_img": the heatmap on nerve and the orginal image
        "pred_tortuosity_level": int value, the predicted tortuosity level
        "pred_levels_probabilities": float values, the predicted probabilities of each level (0, 1, 2, 3)
        "heatmap_on_image": show heatmap on original image
    """

    heatmap_on_image, roi, only_heatmap_on_seg, nerve_gradient_on_img = get_roi(
        img, seg)
    # 'roi': <class 'numpy.ndarray'>
    pred_tortuosity_level, pred_levels_probabilities = grading(img, seg, roi)
    pred_tortuosity_level = pred_tortuosity_level[0]
    pred_levels_probabilities = pred_levels_probabilities[0]
    return heatmap_on_image, only_heatmap_on_seg, nerve_gradient_on_img, pred_tortuosity_level, pred_levels_probabilities


""""-------------我也不知道谁写的，我只是一个搬运工----------"""


class BatchPredictService:
    """BatchPredictService:
        此类封装处理主流程，主要分为流程启动判断，计时，参数过滤/添加, 执行四个步骤：
        notice:
            1. model_path和单例设计模式冲突，综合效率来看决定暂时舍弃Segmentation预测部分的model_path
                - 如果删除单例设计模式，会导致每次运算导入模型从而变慢
                - model_path将会传入工程目录, 然后类中直接组合写死为配置表中的对应模型
                - 此操作会导致后续代码修改难度上升，配置表信息将不能动了
                - 另外父类初始化的模型，将会被自动垃圾回收
    """
    worker = ThreadPool(ApplicationConfig.SystemConfig.NUM_WORKER)
    state_upload_service = MultiStatusUploadImpl()

    def predict_start(self, saver: multibatch.SaveService,
                      save_dir: str, uName: str, token: str, config: dict, **kwargs):

        redis_id = kwargs.get('redis_id')
        pre_message = uName + '-' + token + "-"

        # judge
        if not ApplicationConfig.PathConfig.had_upload(saver.NORMAL_0):
            ConsoleService.console_log(pre_message + 'should upload image first')
            return

        # enhance
        enhance_choice = config['enhance']
        self.upload_task_state(redis_id, "enhance", self.handle_enhance, enhance_choice, saver, pre_message)

        # segmentation
        segmentation_choice = config['segmentation']
        self.upload_task_state(redis_id, "segmentation", self.handle_segmentation, segmentation_choice, saver,
                               pre_message,
                               save_dir=save_dir, enhance_choice=enhance_choice)

        # grade
        grade_choice = config['grade']
        self.upload_task_state(redis_id, "grade", self.handle_grade, grade_choice, saver, pre_message,
                               enhance_choice=enhance_choice)

        # vascular
        vascular_choice = config['vascular']
        vascular_saver = kwargs.get('vascular_saver', None)
        vascular_predictor = kwargs.get('vascular_predictor', None)
        self.upload_task_state(redis_id, "vascular", self.handle_vascular, vascular_choice, saver, pre_message,
                               segmentation_choice=segmentation_choice, enhance_choice=enhance_choice,
                               vascular_saver=vascular_saver, vascular_predictor=vascular_predictor)
        self.state_upload_service.update_total(redis_id, has_finished=True)
        ConsoleService.console_log("Predict Done !!!!!")

    def upload_task_state(self, redis_id, state, func, *args, **kwargs):
        self.state_upload_service.update_state(redis_id, "{} start !!!".format(state))
        r = func(*args, **kwargs, redis_id=redis_id, state_upload=self.state_upload_service)
        self.state_upload_service.update_state(redis_id, "{} done!!!".format(state))
        return r

    def handle_vascular(self, choice: list, saver: multibatch.SaveService, pre_message: str = "", **kwargs):
        vascular_saver = kwargs.get('vascular_saver', None)
        vascular_predictor = kwargs.get('vascular_predictor', None)
        segmentation_choice = kwargs.get('segmentation_choice')
        enhance_choice = kwargs.get('enhance_choice')
        redis_id = kwargs.get("redis_id", None)
        state_upload = kwargs.get('state_upload', None)
        if choice[0] == -1:
            ConsoleService.console_log(pre_message + 'vascular dont start')
        else:
            vascular_predictor.predict_start(saver, vascular_saver, segmentation_choice, pre_message,
                                             redis_id=redis_id, state_upload=state_upload, enhance_choice=enhance_choice)

    def handle_enhance(self, choice: list, saver: multibatch.SaveService, pre_message: str = "", **kwargs):
        if choice[0] == -1:
            ConsoleService.console_log(pre_message + 'enhance dont start')
        else:
            self.start_enhance([saver.NORMAL_0], [saver.ENHANCE_0],
                               os.path.join(ApplicationConfig.PathConfig.SYSTEM_RESOURCE, choice[1]),
                               pre_message, **kwargs)

    def handle_segmentation(self, choice: list, saver: multibatch.SaveService, pre_message: str = "", **kwargs):
        save_dir = kwargs.get('save_dir')
        enhance_choice = kwargs.get('enhance_choice')
        redis_id = kwargs.get('redis_id')

        if choice[0] == -1:
            ConsoleService.console_log(pre_message + 'segmentation dont start')
            return

        from_list, to_list = [saver.NORMAL_0], [saver.NORMAL_1]
        tmp_list = [saver.NORMAL_TMP]
        ceil_list = [saver.NORMAL_1_1]
        if enhance_choice[0] != -1:
            from_list.append(saver.ENHANCE_0)
            to_list.append(saver.ENHANCE_1)
            tmp_list.append(saver.ENHANCE_TMP)
            ceil_list.append(saver.ENHANCE_1_1)

        # model_path = os.path.join(ApplicationConfig.PathConfig.SYSTEM_RESOURCE, choice[1])
        model_path = ApplicationConfig.PathConfig.SYSTEM_RESOURCE
        self.start_segmentation(from_list, to_list, model_path, pre_message,
                                choice=choice[0], redis_id=redis_id, tmp_list=tmp_list, ceil_list=ceil_list)

    def handle_grade(self, choice: list, saver: multibatch.SaveService, pre_message: str = "", **kwargs):
        enhance_choice = kwargs.get('enhance_choice')

        if choice[0] == -1:
            ConsoleService.console_log(pre_message + 'grade dont start')
        else:
            if enhance_choice[0] == -1:
                self.start_grade([(saver.NORMAL_0, saver.NORMAL_1)],
                                 [(saver.NORMAL_2, saver.NORMAL_2_1, saver.NORMAL_3)],
                                 choice[1], pre_message, **kwargs)
            else:
                self.start_grade(
                    [(saver.NORMAL_0, saver.NORMAL_1), (saver.ENHANCE_0, saver.ENHANCE_1)],
                    [(saver.NORMAL_2, saver.NORMAL_2_1, saver.NORMAL_3),
                     (saver.ENHANCE_2, saver.ENHANCE_2_1, saver.ENHANCE_3)],
                    choice[1], pre_message, **kwargs
                )

    @NotifyAnnotation(key=os.path.dirname(__file__), value='grade算法需要并发优化')
    def start_grade(self, from_path: list, to_path: list, model_path: str, pre_message: str, **kwargs):
        for each_from, each_to in zip(from_path, to_path):
            self.start_grade_common(each_from, each_to, pre_message, **kwargs)

    def start_grade_common(self, each_from: str, each_to: str, pre_message: str, **kwargs):
        ori_dir, seg_dir = each_from
        f_2, f_2_1, f_3 = each_to

        # for each_file in os.listdir(ori_dir):
        def for_thread(each_file):
            abs_ori_path = os.path.join(ori_dir, each_file)
            abs_seg_path = os.path.join(seg_dir, each_file)

            seg = cv2.imread(abs_seg_path, 0)
            seg = cv2.resize(seg, (384, 384))
            ret, seg = cv2.threshold(seg, 127, 255, cv2.THRESH_BINARY)

            img = Image.open(abs_ori_path)
            # org_size = img.size
            img = img.convert('L')
            img = img.resize((384, 384))
            img_cv = cv2.cvtColor(np.asarray(img), cv2.IMREAD_GRAYSCALE)[:, :, 2]
            img_cv = cv2.resize(img_cv, (384, 384))

            # 热力图
            heatmap_on_image, only_heatmap_on_seg, nerve_gradient_on_img, pred_tortuosity_level, pred_levels_probabilities \
                = tortuosity_grading(img_cv, seg)

            # nerve_gradient_on_img = nerve_gradient_on_img.resize(org_size)

            abs_save_f_2 = os.path.join(f_2, each_file)
            nerve_gradient_on_img.save(abs_save_f_2)

            abs_save_f_2_1 = os.path.join(f_2_1, each_file)
            with open(abs_save_f_2_1 + ".txt", 'w') as f:
                level = int(pred_tortuosity_level)
                d = {
                    'probabilities': float(list(pred_levels_probabilities)[level]),
                    'level': level + 1
                }
                f.write(json.dumps(d))

            # heatmap_on_image = heatmap_on_image.resize(org_size)
            abs_save_f3 = os.path.join(f_3, each_file)
            heatmap_on_image.save(abs_save_f3)

            IStatusUpload.update_state_from_kwargs(**kwargs, token=each_file)
            ConsoleService.console_log(pre_message + abs_save_f3 + " grade量化成功")

        BatchPredictService.worker.map(for_thread, os.listdir(ori_dir))

    def start_segmentation(self, from_path: list, to_path: list, model_path: str, pre_message: str, **kwargs):
        choice = kwargs.get('choice')
        if choice == ApplicationConfig.NetConfig.SEGMENTATION_USE_CSNET:
            csnet_batch_factory = CSNetBatchFactory()
            csnet_predict_service = CSNetService(model_path)
            IPredictHandler.predict(from_path, to_path,
                                    csnet_batch_factory,
                                    csnet_predict_service,
                                    ApplicationConfig.SystemConfig.BATCH_SIZE_ENHANCE,
                                    "{}csnet分割".format(pre_message), **kwargs)
        elif choice == ApplicationConfig.NetConfig.SEGMENTATION_USE_NERVE:
            nerve_batch_factory = NerveFormerBatchFactory()
            nerve_predict_service = NerveFormerService(model_path)
            IPredictHandler.predict(from_path, to_path,
                                    nerve_batch_factory,
                                    nerve_predict_service,
                                    ApplicationConfig.SystemConfig.BATCH_SIZE_ENHANCE,
                                    "{}nerve分割".format(pre_message), **kwargs)
        elif choice == ApplicationConfig.NetConfig.SEGMENTATION_USE_CSNET_CEIL:
            csnet_ceil_batch_factory = CSNetCeilBatchFactory()
            csnet_ceil_predict_service = CSNetCeilService(model_path)
            IPredictHandler.predict(from_path, to_path,
                                    csnet_ceil_batch_factory,
                                    csnet_ceil_predict_service,
                                    ApplicationConfig.SystemConfig.BATCH_SIZE_ENHANCE,
                                    "{}csnet ceil分割".format(pre_message), **kwargs)
        elif choice == ApplicationConfig.NetConfig.SEGMENTATION_USE_NERVE_CEIL:
            nerve_former_ceil_batch_factory = NerveCeilBatchFactory()
            nerve_former_ceil_predict_factory = NerveFormerCeilService(model_path)
            IPredictHandler.predict(from_path, to_path,
                                    nerve_former_ceil_batch_factory,
                                    nerve_former_ceil_predict_factory,
                                    ApplicationConfig.SystemConfig.BATCH_SIZE_ENHANCE,
                                    "{}nerve_former ceil分割".format(pre_message), **kwargs)

    def start_enhance(self, from_path: list, to_path: list, model_path: str, pre_message: str, **kwargs):
        stillgan_batch_factory = StillganBatchFactory()
        stillgan_predict_service = StillganService(model_path)
        IPredictHandler.predict(from_path, to_path,
                                stillgan_batch_factory,
                                stillgan_predict_service,
                                ApplicationConfig.SystemConfig.BATCH_SIZE_ENHANCE,
                                "{}stillgan增强完成".format(pre_message),
                                **kwargs)
